import os
import torch
from PIL import Image
from torchvision import transforms
import time

INPUT_IMAGE_PATH = r"F:\project\jili_three_project\foggy_project_all\1108_new\two\1111\000755.jpg"  # 000755
OUTPUT_DIR = r"F:\project\jili_three_project\foggy_project_all\1108_new\two\aug_check"

IMAGE_SIZE = (360, 640)
NORMALIZE_MEAN = [0.485, 0.456, 0.406]
NORMALIZE_STD = [0.229, 0.224, 0.225]
train_transform = transforms.Compose([
    transforms.Resize(IMAGE_SIZE),
    transforms.RandomApply([transforms.ColorJitter(brightness=0.15, contrast=0.15, saturation=0.1)], p=0.5),
    transforms.RandomApply([transforms.GaussianBlur(kernel_size=3, sigma=(0.1, 0.5))], p=0.3),
    transforms.ToTensor(),
    transforms.RandomErasing(p=0.5, scale=(0.02, 0.1), ratio=(0.3, 3.3), value='random'),
    transforms.Normalize(mean=NORMALIZE_MEAN, std=NORMALIZE_STD),
])


def denormalize_and_to_pil(tensor):
    """
    将一个归一化后的 [C, H, W] tensor 转换回 PIL Image 以便查看。
    """
    img_tensor = tensor.clone().cpu()

    mean = torch.tensor(NORMALIZE_MEAN).view(3, 1, 1)
    std = torch.tensor(NORMALIZE_STD).view(3, 1, 1)

    img_tensor = img_tensor * std + mean

    img_tensor = torch.clamp(img_tensor, 0, 1)

    pil_image = transforms.ToPILImage()(img_tensor)

    return pil_image


def visualize_augmentation():
    os.makedirs(OUTPUT_DIR, exist_ok=True)

    if not os.path.exists(INPUT_IMAGE_PATH):
        print(f"** 错误: 找不到输入图片: {INPUT_IMAGE_PATH} **")
        print("请修改脚本第 6 行的 INPUT_IMAGE_PATH 变量。")
        return

    try:
        original_image = Image.open(INPUT_IMAGE_PATH).convert("RGB")
    except Exception as e:
        print(f"** 错误: 加载图片失败: {e} **")
        return

    print(f"加载原始图片: {INPUT_IMAGE_PATH}")
    print(f"应用随机数据增强 (模糊, 颜色抖动)...")

    original_resized = original_image.resize(IMAGE_SIZE[::-1])  # PIL.Image.resize 接收 (W, H)
    original_resized.save(os.path.join(OUTPUT_DIR, "0_original_resized.jpg"))

    num_examples = 10
    print(f"正在生成 {num_examples} 个增强后的样本...")

    for i in range(num_examples):
        augmented_tensor = train_transform(original_image)

        augmented_pil_image = denormalize_and_to_pil(augmented_tensor)

        timestamp = int(time.time() * 1000)  # 用时间戳防止重名
        output_path = os.path.join(OUTPUT_DIR, f"augmented_example_{i + 1}.jpg")
        augmented_pil_image.save(output_path)

    print(f"\n成功!")
    print(f"请检查文件夹: {OUTPUT_DIR}")
    print(f"那里现在应该有 1 张 'original_resized.jpg' (用于对比)")
    print(f"以及 {num_examples} 张 'augmented_example_*.jpg' (增强后的结果)")


if __name__ == "__main__":
    visualize_augmentation()